{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Part 1, Topic 1, Lab B: AES Loop Skip Fault Attack in Practice"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "NOTE: This lab references some (commercial) training material on [ChipWhisperer.io](https://www.ChipWhisperer.io). You can freely execute and use the lab per the open-source license (including using it in your own courses if you distribute similarly), but you must maintain notice about this source location. Consider joining our training course to enjoy the full experience.\n",
    "\n",
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**SUMMARY:** *In the last lab, we showed recovering an AES becomes trivial if we're able to skip the repeated rounds of AES. In this lab, we'll look at applying this attack to a real AES implementation.*\n",
    "\n",
    "**This lab is only supported with TINYAES128C firmware**\n",
    "\n",
    "**LEARNING OUTCOMES:**\n",
    "* Understanding how C code can be modified by a compiler\n",
    "* Using a theoretical fault model to mount a real attack\n",
    "\n",
    "## Can we skip the loop of AES\n",
    "\n",
    "For this particular fault, the answer is that it depends! Let's take a look at the source code for TINYAES:\n",
    "\n",
    "```C\n",
    "// Cipher is the main function that encrypts the PlainText.\n",
    "static void Cipher(void)\n",
    "{\n",
    "  uint8_t round = 0;\n",
    "\n",
    "  // Add the First round key to the state before starting the rounds.\n",
    "  AddRoundKey(0); \n",
    "  \n",
    "  // There will be Nr rounds.\n",
    "  // The first Nr-1 rounds are identical.\n",
    "  // These Nr-1 rounds are executed in the loop below.\n",
    "  for(round = 1; round < Nr; ++round)\n",
    "  {\n",
    "    SubBytes();\n",
    "    ShiftRows();\n",
    "    MixColumns();\n",
    "    AddRoundKey(round);\n",
    "  }\n",
    "  \n",
    "  // The last round is given below.\n",
    "  // The MixColumns function is not here in the last round.\n",
    "  SubBytes();\n",
    "  ShiftRows();\n",
    "  AddRoundKey(Nr);\n",
    "}\n",
    "```\n",
    "\n",
    "We learned in Fault101 that fault injection can be used to bypass the final check and stay in a loop longer. Similarly, glitching can also be used to break out of a loop early. This is likely one of the effects that we saw when glitching past a password check in that course as well. That being said, it's up to the compiler how this loop gets translated to assembly. For example, `Nr` is constant here, so the compiler might decided to completely unroll the loop, preventing us from ever breaking free! For our glitch to work we really have three requirements:\n",
    "\n",
    "1. There actually needs to be a loop for us to break free from. The compiler can't unroll the loop.\n",
    "1. The registers need to line up properly such that, if we skip the loop, the last round still completes properly.\n",
    "1. The compiler can't skip the first `round < Nr` check. If the compiler doesn't know that `round` always starts off less than `Nr`, it needs to do a check at the start of the loop.\n",
    "\n",
    "unfortunately for us, while our default build of TINYAES128C will fulfill the first two requirements, if we look at the generated assembly from `simpleserial-aes-PLATFORM.lss` (CWLITEARM in this case):\n",
    "\n",
    "```C\n",
    "// Cipher is the main function that encrypts the PlainText.\n",
    "static void Cipher(void)\n",
    "{\n",
    " 8001378:\te92d 4ff8 \tstmdb\tsp!, {r3, r4, r5, r6, r7, r8, r9, sl, fp, lr}\n",
    "  uint8_t round = 0;\n",
    "\n",
    "  // Add the First round key to the state before starting the rounds.\n",
    "  AddRoundKey(0); \n",
    " 800137c:\t2000      \tmovs\tr0, #0\n",
    " 800137e:\tf7ff ffa1 \tbl\t80012c4 <AddRoundKey>\n",
    "  \n",
    "  // There will be Nr rounds.\n",
    "  // The first Nr-1 rounds are identical.\n",
    "  // These Nr-1 rounds are executed in the loop below.\n",
    "  for(round = 1; round < Nr; ++round)\n",
    " 8001382:\t2401      \tmovs\tr4, #1\n",
    "  {\n",
    "    SubBytes();\n",
    " 8001384:\tf7ff ffb8 \tbl\t80012f8 <SubBytes>\n",
    "    ShiftRows();\n",
    " 8001388:\tf7ff ffce \tbl\t8001328 <ShiftRows>\n",
    "  for(i = 0; i < 4; ++i)\n",
    " 800138c:\t4b1f      \tldr\tr3, [pc, #124]\t; (800140c <Cipher+0x94>)\n",
    " 800138e:\tf8d3 10b4 \tldr.w\tr1, [r3, #180]\t; 0xb4\n",
    " 8001392:\tf101 0b10 \tadd.w\tfp, r1, #16\n",
    "    t   = (*state)[i][0];\n",
    " 8001396:\tf891 a000 \tldrb.w\tsl, [r1]\n",
    "    Tmp = (*state)[i][0] ^ (*state)[i][1] ^ (*state)[i][2] ^ (*state)[i][3] ;\n",
    " 800139a:\t784e      \tldrb\tr6, [r1, #1]\n",
    " 800139c:\t788d      \tldrb\tr5, [r1, #2]\n",
    " 800139e:\tf891 9003 \tldrb.w\tr9, [r1, #3]\n",
    " 80013a2:\tea8a 0006 \teor.w\tr0, sl, r6\n",
    " 80013a6:\tea85 0809 \teor.w\tr8, r5, r9\n",
    " 80013aa:\tea88 0700 \teor.w\tr7, r8, r0\n",
    "    Tm  = (*state)[i][0] ^ (*state)[i][1] ; Tm = xtime(Tm);  (*state)[i][0] ^= Tm ^ Tmp ;\n",
    " 80013ae:\tf7ff ffd9 \tbl\t8001364 <xtime>\n",
    " 80013b2:\tea8a 0000 \teor.w\tr0, sl, r0\n",
    " 80013b6:\t4078      \teors\tr0, r7\n",
    " 80013b8:\t7008      \tstrb\tr0, [r1, #0]\n",
    "    Tm  = (*state)[i][1] ^ (*state)[i][2] ; Tm = xtime(Tm);  (*state)[i][1] ^= Tm ^ Tmp ;\n",
    " 80013ba:\tea86 0005 \teor.w\tr0, r6, r5\n",
    " 80013be:\tf7ff ffd1 \tbl\t8001364 <xtime>\n",
    " 80013c2:\t4046      \teors\tr6, r0\n",
    " 80013c4:\t407e      \teors\tr6, r7\n",
    " 80013c6:\t704e      \tstrb\tr6, [r1, #1]\n",
    "    Tm  = (*state)[i][2] ^ (*state)[i][3] ; Tm = xtime(Tm);  (*state)[i][2] ^= Tm ^ Tmp ;\n",
    " 80013c8:\t4640      \tmov\tr0, r8\n",
    " 80013ca:\tf7ff ffcb \tbl\t8001364 <xtime>\n",
    " 80013ce:\t4045      \teors\tr5, r0\n",
    " 80013d0:\t407d      \teors\tr5, r7\n",
    " 80013d2:\t708d      \tstrb\tr5, [r1, #2]\n",
    "    Tm  = (*state)[i][3] ^ t ;        Tm = xtime(Tm);  (*state)[i][3] ^= Tm ^ Tmp ;\n",
    " 80013d4:\tea8a 0009 \teor.w\tr0, sl, r9\n",
    " 80013d8:\tf7ff ffc4 \tbl\t8001364 <xtime>\n",
    " 80013dc:\tea89 0900 \teor.w\tr9, r9, r0\n",
    " 80013e0:\tea87 0709 \teor.w\tr7, r7, r9\n",
    " 80013e4:\t70cf      \tstrb\tr7, [r1, #3]\n",
    "  for(i = 0; i < 4; ++i)\n",
    " 80013e6:\t3104      \tadds\tr1, #4\n",
    " 80013e8:\t4559      \tcmp\tr1, fp\n",
    " 80013ea:\td1d4      \tbne.n\t8001396 <Cipher+0x1e>\n",
    "    MixColumns();\n",
    "    AddRoundKey(round);\n",
    " 80013ec:\t4620      \tmov\tr0, r4\n",
    "  for(round = 1; round < Nr; ++round)\n",
    " 80013ee:\t3401      \tadds\tr4, #1\n",
    " 80013f0:\tb2e4      \tuxtb\tr4, r4\n",
    "    AddRoundKey(round);\n",
    " 80013f2:\tf7ff ff67 \tbl\t80012c4 <AddRoundKey>\n",
    "  for(round = 1; round < Nr; ++round)\n",
    " 80013f6:\t2c0a      \tcmp\tr4, #10                ; <-- round != Nr check\n",
    " 80013f8:\td1c4      \tbne.n\t8001384 <Cipher+0xc> ; <--\n",
    "  }\n",
    "  \n",
    "  // The last round is given below.\n",
    "  // The MixColumns function is not here in the last round.\n",
    "  SubBytes();\n",
    " 80013fa:\tf7ff ff7d \tbl\t80012f8 <SubBytes>\n",
    "  ShiftRows();\n",
    " 80013fe:\tf7ff ff93 \tbl\t8001328 <ShiftRows>\n",
    "  AddRoundKey(Nr);\n",
    " 8001402:\t4620      \tmov\tr0, r4\n",
    "}\n",
    " 8001404:\te8bd 4ff8 \tldmia.w\tsp!, {r3, r4, r5, r6, r7, r8, r9, sl, fp, lr}\n",
    "\n",
    "```\n",
    "it won't fulfill the last one. This means that the device will always go through one round of AES before we have the opportunity to break out of the loop. An attack is still possible on this single round of AES (with only 2 faults to boot), but the math behind it is much more complicated. We'll look at this attack in the next lab, but it would still be good to see our simpler attack working on real hardware. Luckily, it's actually pretty easy to get the firmware to fulfill the last  requirement! All we need to do is make `round` a volatile variable:\n",
    "\n",
    "```C\n",
    "// Cipher is the main function that encrypts the PlainText.\n",
    "static void Cipher(void)\n",
    "{\n",
    "  volatile uint8_t round = 0;\n",
    "\n",
    "  // Add the First round key to the state before starting the rounds.\n",
    "  AddRoundKey(0); \n",
    "  \n",
    "  // There will be Nr rounds.\n",
    "  // The first Nr-1 rounds are identical.\n",
    "  // These Nr-1 rounds are executed in the loop below.\n",
    "  for(round = 1; round < Nr; ++round)\n",
    "  {\n",
    "    SubBytes();\n",
    "    ShiftRows();\n",
    "    MixColumns();\n",
    "    AddRoundKey(round);\n",
    "  }\n",
    "  \n",
    "  // The last round is given below.\n",
    "  // The MixColumns function is not here in the last round.\n",
    "  SubBytes();\n",
    "  ShiftRows();\n",
    "  AddRoundKey(Nr);\n",
    "}\n",
    "```\n",
    "\n",
    "This will result in the following assembly (again on CWLITEARM):\n",
    "\n",
    "```C\n",
    "// Cipher is the main function that encrypts the PlainText.\n",
    "static void Cipher(void)\n",
    "{\n",
    " 8001378:\te92d 4ff7 \tstmdb\tsp!, {r0, r1, r2, r4, r5, r6, r7, r8, r9, sl, fp, lr}\n",
    "  volatile uint8_t round = 0;\n",
    " 800137c:\t2000      \tmovs\tr0, #0\n",
    " 800137e:\tf88d 0007 \tstrb.w\tr0, [sp, #7]\n",
    "    t   = (*state)[i][0];\n",
    " 8001382:\t4f29      \tldr\tr7, [pc, #164]\t; (8001428 <Cipher+0xb0>)\n",
    "\n",
    "  // Add the First round key to the state before starting the rounds.\n",
    "  AddRoundKey(0); \n",
    " 8001384:\tf7ff ff9e \tbl\t80012c4 <AddRoundKey>\n",
    "  \n",
    "  // There will be Nr rounds.\n",
    "  // The first Nr-1 rounds are identical.\n",
    "  // These Nr-1 rounds are executed in the loop below.\n",
    "  for(round = 1; round < Nr; ++round)\n",
    " 8001388:\t2301      \tmovs\tr3, #1\n",
    " 800138a:\tf88d 3007 \tstrb.w\tr3, [sp, #7]\n",
    " 800138e:\tf89d 3007 \tldrb.w\tr3, [sp, #7]\n",
    " 8001392:\t2b09      \tcmp\tr3, #9                  ; <-- round < Nr check\n",
    " 8001394:\td909      \tbls.n\t80013aa <Cipher+0x32> ; <--\n",
    "    AddRoundKey(round);\n",
    "  }\n",
    "  \n",
    "  // The last round is given below.\n",
    "  // The MixColumns function is not here in the last round.\n",
    "  SubBytes();\n",
    " 8001396:\tf7ff ffaf \tbl\t80012f8 <SubBytes>\n",
    "  ShiftRows();\n",
    " 800139a:\tf7ff ffc5 \tbl\t8001328 <ShiftRows>\n",
    "  AddRoundKey(Nr);\n",
    " 800139e:\t200a      \tmovs\tr0, #10\n",
    "}\n",
    " 80013a0:\tb003      \tadd\tsp, #12\n",
    " 80013a2:\te8bd 4ff0 \tldmia.w\tsp!, {r4, r5, r6, r7, r8, r9, sl, fp, lr}\n",
    "  AddRoundKey(Nr);\n",
    " 80013a6:\tf7ff bf8d \tb.w\t80012c4 <AddRoundKey>\n",
    "    SubBytes();\n",
    " 80013aa:\tf7ff ffa5 \tbl\t80012f8 <SubBytes>\n",
    "    ShiftRows();\n",
    " 80013ae:\tf7ff ffbb \tbl\t8001328 <ShiftRows>\n",
    "  for(i = 0; i < 4; ++i)\n",
    " 80013b2:\tf8d7 10b4 \tldr.w\tr1, [r7, #180]\t; 0xb4\n",
    " 80013b6:\tf101 0a10 \tadd.w\tsl, r1, #16\n",
    "    t   = (*state)[i][0];\n",
    " 80013ba:\tf891 9000 \tldrb.w\tr9, [r1]\n",
    "    Tmp = (*state)[i][0] ^ (*state)[i][1] ^ (*state)[i][2] ^ (*state)[i][3] ;\n",
    " 80013be:\t784d      \tldrb\tr5, [r1, #1]\n",
    " 80013c0:\t788c      \tldrb\tr4, [r1, #2]\n",
    " 80013c2:\tf891 8003 \tldrb.w\tr8, [r1, #3]\n",
    " 80013c6:\tea89 0005 \teor.w\tr0, r9, r5\n",
    " 80013ca:\tea84 0b08 \teor.w\tfp, r4, r8\n",
    " 80013ce:\tea8b 0600 \teor.w\tr6, fp, r0\n",
    "    Tm  = (*state)[i][0] ^ (*state)[i][1] ; Tm = xtime(Tm);  (*state)[i][0] ^= Tm ^ Tmp ;\n",
    " 80013d2:\tf7ff ffc7 \tbl\t8001364 <xtime>\n",
    " 80013d6:\tea89 0000 \teor.w\tr0, r9, r0\n",
    " 80013da:\t4070      \teors\tr0, r6\n",
    " 80013dc:\t7008      \tstrb\tr0, [r1, #0]\n",
    "    Tm  = (*state)[i][1] ^ (*state)[i][2] ; Tm = xtime(Tm);  (*state)[i][1] ^= Tm ^ Tmp ;\n",
    " 80013de:\tea85 0004 \teor.w\tr0, r5, r4\n",
    " 80013e2:\tf7ff ffbf \tbl\t8001364 <xtime>\n",
    " 80013e6:\t4045      \teors\tr5, r0\n",
    " 80013e8:\t4075      \teors\tr5, r6\n",
    " 80013ea:\t704d      \tstrb\tr5, [r1, #1]\n",
    "    Tm  = (*state)[i][2] ^ (*state)[i][3] ; Tm = xtime(Tm);  (*state)[i][2] ^= Tm ^ Tmp ;\n",
    " 80013ec:\t4658      \tmov\tr0, fp\n",
    " 80013ee:\tf7ff ffb9 \tbl\t8001364 <xtime>\n",
    " 80013f2:\t4044      \teors\tr4, r0\n",
    " 80013f4:\t4074      \teors\tr4, r6\n",
    " 80013f6:\t708c      \tstrb\tr4, [r1, #2]\n",
    "    Tm  = (*state)[i][3] ^ t ;        Tm = xtime(Tm);  (*state)[i][3] ^= Tm ^ Tmp ;\n",
    " 80013f8:\tea89 0008 \teor.w\tr0, r9, r8\n",
    " 80013fc:\tf7ff ffb2 \tbl\t8001364 <xtime>\n",
    " 8001400:\tea88 0800 \teor.w\tr8, r8, r0\n",
    " 8001404:\tea86 0608 \teor.w\tr6, r6, r8\n",
    " 8001408:\t70ce      \tstrb\tr6, [r1, #3]\n",
    "  for(i = 0; i < 4; ++i)\n",
    " 800140a:\t3104      \tadds\tr1, #4\n",
    " 800140c:\t4551      \tcmp\tr1, sl\n",
    " 800140e:\td1d4      \tbne.n\t80013ba <Cipher+0x42>\n",
    "    AddRoundKey(round);\n",
    " 8001410:\tf89d 0007 \tldrb.w\tr0, [sp, #7]\n",
    " 8001414:\tf7ff ff56 \tbl\t80012c4 <AddRoundKey>\n",
    "  for(round = 1; round < Nr; ++round)\n",
    " 8001418:\tf89d 3007 \tldrb.w\tr3, [sp, #7]\n",
    " 800141c:\t3301      \tadds\tr3, #1\n",
    " 800141e:\tb2db      \tuxtb\tr3, r3\n",
    " 8001420:\tf88d 3007 \tstrb.w\tr3, [sp, #7]\n",
    " 8001424:\te7b3      \tb.n\t800138e <Cipher+0x16>\n",
    " 8001426:\tbf00      \tnop\n",
    " 8001428:\t200006dc \t.word\t0x200006dc\n",
    "```\n",
    "\n",
    "This time, we can see the check is being done at the beginning of the loop instead of at the end at address 0x8001392. If you look closely, you can see if we skip this check, we go directly into the last round! Make this change in your own `firmware/mcu/crypto/TINYAES128C/aes.c` and we can start the lab!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "SCOPETYPE = 'OPENADC'\n",
    "PLATFORM = 'CWLITEARM'\n",
    "CRYPTO_TARGET = 'TINYAES128C'\n",
    "SS_VER='SS_VER_2_1'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%run \"../../Setup_Scripts/Setup_Generic.ipynb\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%bash -s \"$PLATFORM\" \"$CRYPTO_TARGET\" \"$SS_VER\"\n",
    "cd ../../../firmware/mcu/simpleserial-aes\n",
    "make PLATFORM=$1 CRYPTO_TARGET=$2 SS_VER=$3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fw_path = \"../../../firmware/mcu/simpleserial-aes/simpleserial-aes-{}.hex\".format(PLATFORM)\n",
    "cw.program_target(scope, prog, fw_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if PLATFORM == \"CWLITEXMEGA\":\n",
    "    def reboot_flush():            \n",
    "        scope.io.pdic = False\n",
    "        time.sleep(0.1)\n",
    "        scope.io.pdic = \"high_z\"\n",
    "        time.sleep(0.1)\n",
    "        #Flush garbage too\n",
    "        target.flush()\n",
    "else:\n",
    "    def reboot_flush():            \n",
    "        scope.io.nrst = False\n",
    "        time.sleep(0.05)\n",
    "        scope.io.nrst = \"high_z\"\n",
    "        time.sleep(0.05)\n",
    "        #Flush garbage too\n",
    "        target.flush()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Getting a Fault\n",
    "\n",
    "There's two problems that we're going to face in getting a fault out of this code:\n",
    "\n",
    "1. Where do we insert the fault? Obviously, it'll be somewhere near the beginning, but it would be nice to narrow down the location\n",
    "1. How do we know we've gotten the specific fault we're looking for? Inserting faults nearby to the correct location will generate faults in the output, but it's hard to tell just from looking at the output if we've broken out of the loop.\n",
    "\n",
    "We can pretty easily solve both these problems with power analysis! Since the different operations of AES are pretty distinct in AES, we can visually inspect an unfaulted power trace to know where to insert the fault. For the second issue, we can again visually inspect the power trace. Breaking out of the loop will look very different to completing the rest of AES.\n",
    "\n",
    "Let's get the target to encrypt something and capture a power trace. We'll also capture a copy of the ciphertext, which we can use to detect glitches in general:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "scope.clock.adc_src = \"clkgen_x1\"\n",
    "reboot_flush()\n",
    "scope.arm()\n",
    "target.simpleserial_write('p', bytearray([0]*16))\n",
    "ret = scope.capture()\n",
    "if ret:\n",
    "    print(\"No trigger!\")\n",
    "\n",
    "wave = scope.get_last_trace()\n",
    "\n",
    "output = target.simpleserial_read_witherrors('r', 16)\n",
    "gold_ct = output['payload']\n",
    "\n",
    "print(gold_ct)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cw.plot(wave[:2000])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now select a range of values of where to glitch. Remember the add round key operation is done at the beginning, then at the end of very loop through an AES round.  You should be glitching between the last little bit of this operation and the beginning of the next one."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "glitch_loc = range(300, 340)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "By this point, you should have some pretty reliable settings for glitching the target, so the provided glitch loop won't even use the GlitchController, but feel free to add it in if you're not super confident in using only one pair of glitch settings."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if scope._is_husky:\n",
    "    scope.glitch.enabled = True\n",
    "scope.glitch.clk_src = \"clkgen\"\n",
    "scope.glitch.output = \"clock_xor\"\n",
    "scope.glitch.trigger_src = \"ext_single\"\n",
    "scope.glitch.repeat = 1\n",
    "scope.io.hs2 = \"glitch\"\n",
    "# These width/offset numbers are for CW-Lite/Pro; for CW-Husky, convert as per Fault 1_1:\n",
    "scope.glitch.width = 3\n",
    "scope.glitch.offset = -12.8\n",
    "print(scope.glitch)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You could use a SAD comparison to detect when the glitch looks substantially different, but it'll be a bit easier to loop until we glitch, then check what the waveform looks like. If you get a glitch, but it doesn't result in the correct effect, you can pretty safely skip that glitch spot."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from tqdm.notebook import tqdm, trange\n",
    "wave = None\n",
    "import logging\n",
    "ktp = cw.ktp.Basic()\n",
    "logging.getLogger().setLevel(logging.ERROR)\n",
    "reboot_flush()\n",
    "for i in trange(min(glitch_loc), max(glitch_loc) + 1):\n",
    "    scope.adc.timeout = 0.2\n",
    "    scope.glitch.ext_offset = i\n",
    "    ack = None\n",
    "    while ack is None:\n",
    "        target.simpleserial_write('k', ktp.next()[0])\n",
    "        ack = target.simpleserial_wait_ack()\n",
    "        if ack is None:\n",
    "            reboot_flush()\n",
    "            time.sleep(0.1)\n",
    "    \n",
    "    scope.arm()\n",
    "    \n",
    "    pt = bytearray([0]*16)\n",
    "    target.simpleserial_write('p', pt)\n",
    "    ret = scope.capture()\n",
    "    if ret:\n",
    "        reboot_flush() #bad if we accidentally didn't have this work\n",
    "        time.sleep(0.1)\n",
    "        print(\"timed out!\")\n",
    "        continue\n",
    "    output = target.simpleserial_read_witherrors('r', 16, glitch_timeout = 1)\n",
    "    if output['valid']:\n",
    "        if output['payload'] != gold_ct:\n",
    "            print(\"Glitched at {}\".format(i))\n",
    "            wave = scope.get_last_trace()\n",
    "            break\n",
    "    else:\n",
    "        reboot_flush()\n",
    "        \n",
    "cw.plot(wave)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's collect our glitched ciphertext:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "glitched_ct0 = bytearray(output['payload'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "glitched_ct0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pt0 = bytearray(pt)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can repeat this twice more with different plaintexts to get all the faults needed:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from tqdm.notebook import tqdm, trange\n",
    "wave = None\n",
    "import logging\n",
    "ktp = cw.ktp.Basic()\n",
    "logging.getLogger().setLevel(logging.ERROR)\n",
    "reboot_flush()\n",
    "while True:\n",
    "    scope.adc.timeout = 0.2\n",
    "    scope.glitch.ext_offset = i #should still be in the right spot from the last glitch\n",
    "    ack = None\n",
    "    while ack is None:\n",
    "        target.simpleserial_write('k', ktp.next()[0])\n",
    "        ack = target.simpleserial_wait_ack()\n",
    "        if ack is None:\n",
    "            reboot_flush()\n",
    "            time.sleep(0.1)\n",
    "    \n",
    "    scope.arm()\n",
    "    \n",
    "    pt = bytearray([1]*16)\n",
    "    target.simpleserial_write('p', pt)\n",
    "    ret = scope.capture()\n",
    "    if ret:\n",
    "        reboot_flush() #bad if we accidentally didn't have this work\n",
    "        time.sleep(0.1)\n",
    "        print(\"timed out!\")\n",
    "        continue\n",
    "    output = target.simpleserial_read_witherrors('r', 16, glitch_timeout = 1)\n",
    "    if output['valid']:\n",
    "        if output['payload'] != gold_ct:\n",
    "            print(\"Glitched at {}\".format(i))\n",
    "            wave = scope.get_last_trace()\n",
    "            break\n",
    "    else:\n",
    "        reboot_flush()\n",
    "        \n",
    "cw.plot(wave)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "glitched_ct1 = bytearray(output['payload'])\n",
    "pt1 = bytearray(pt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from tqdm.notebook import tqdm, trange\n",
    "wave = None\n",
    "import logging\n",
    "ktp = cw.ktp.Basic()\n",
    "logging.getLogger().setLevel(logging.ERROR)\n",
    "reboot_flush()\n",
    "while True:\n",
    "    scope.adc.timeout = 0.2\n",
    "    scope.glitch.ext_offset = i #should still be in the right spot from the last glitch\n",
    "    ack = None\n",
    "    while ack is None:\n",
    "        target.simpleserial_write('k', ktp.next()[0])\n",
    "        ack = target.simpleserial_wait_ack()\n",
    "        if ack is None:\n",
    "            reboot_flush()\n",
    "            time.sleep(0.1)\n",
    "    \n",
    "    scope.arm()\n",
    "    \n",
    "    pt = bytearray([2]*16)\n",
    "    target.simpleserial_write('p', pt)\n",
    "    ret = scope.capture()\n",
    "    if ret:\n",
    "        reboot_flush() #bad if we accidentally didn't have this work\n",
    "        time.sleep(0.1)\n",
    "        print(\"timed out!\")\n",
    "        continue\n",
    "    output = target.simpleserial_read_witherrors('r', 16, glitch_timeout = 1)\n",
    "    if output['valid']:\n",
    "        if output['payload'] != gold_ct:\n",
    "            print(\"Glitched at {}\".format(i))\n",
    "            wave = scope.get_last_trace()\n",
    "            break\n",
    "    else:\n",
    "        reboot_flush()\n",
    "        \n",
    "cw.plot(wave)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "glitched_ct2 = bytearray(output['payload'])\n",
    "pt2 = bytearray(pt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sbox = [\n",
    "    # 0    1    2    3    4    5    6    7    8    9    a    b    c    d    e    f \n",
    "    0x63,0x7c,0x77,0x7b,0xf2,0x6b,0x6f,0xc5,0x30,0x01,0x67,0x2b,0xfe,0xd7,0xab,0x76, # 0\n",
    "    0xca,0x82,0xc9,0x7d,0xfa,0x59,0x47,0xf0,0xad,0xd4,0xa2,0xaf,0x9c,0xa4,0x72,0xc0, # 1\n",
    "    0xb7,0xfd,0x93,0x26,0x36,0x3f,0xf7,0xcc,0x34,0xa5,0xe5,0xf1,0x71,0xd8,0x31,0x15, # 2\n",
    "    0x04,0xc7,0x23,0xc3,0x18,0x96,0x05,0x9a,0x07,0x12,0x80,0xe2,0xeb,0x27,0xb2,0x75, # 3\n",
    "    0x09,0x83,0x2c,0x1a,0x1b,0x6e,0x5a,0xa0,0x52,0x3b,0xd6,0xb3,0x29,0xe3,0x2f,0x84, # 4\n",
    "    0x53,0xd1,0x00,0xed,0x20,0xfc,0xb1,0x5b,0x6a,0xcb,0xbe,0x39,0x4a,0x4c,0x58,0xcf, # 5\n",
    "    0xd0,0xef,0xaa,0xfb,0x43,0x4d,0x33,0x85,0x45,0xf9,0x02,0x7f,0x50,0x3c,0x9f,0xa8, # 6\n",
    "    0x51,0xa3,0x40,0x8f,0x92,0x9d,0x38,0xf5,0xbc,0xb6,0xda,0x21,0x10,0xff,0xf3,0xd2, # 7\n",
    "    0xcd,0x0c,0x13,0xec,0x5f,0x97,0x44,0x17,0xc4,0xa7,0x7e,0x3d,0x64,0x5d,0x19,0x73, # 8\n",
    "    0x60,0x81,0x4f,0xdc,0x22,0x2a,0x90,0x88,0x46,0xee,0xb8,0x14,0xde,0x5e,0x0b,0xdb, # 9\n",
    "    0xe0,0x32,0x3a,0x0a,0x49,0x06,0x24,0x5c,0xc2,0xd3,0xac,0x62,0x91,0x95,0xe4,0x79, # a\n",
    "    0xe7,0xc8,0x37,0x6d,0x8d,0xd5,0x4e,0xa9,0x6c,0x56,0xf4,0xea,0x65,0x7a,0xae,0x08, # b\n",
    "    0xba,0x78,0x25,0x2e,0x1c,0xa6,0xb4,0xc6,0xe8,0xdd,0x74,0x1f,0x4b,0xbd,0x8b,0x8a, # c\n",
    "    0x70,0x3e,0xb5,0x66,0x48,0x03,0xf6,0x0e,0x61,0x35,0x57,0xb9,0x86,0xc1,0x1d,0x9e, # d\n",
    "    0xe1,0xf8,0x98,0x11,0x69,0xd9,0x8e,0x94,0x9b,0x1e,0x87,0xe9,0xce,0x55,0x28,0xdf, # e\n",
    "    0x8c,0xa1,0x89,0x0d,0xbf,0xe6,0x42,0x68,0x41,0x99,0x2d,0x0f,0xb0,0x54,0xbb,0x16  # f\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "key_guess = []\n",
    "for kbyte in range(16):\n",
    "    for i in range(255):\n",
    "        if (sbox[i ^ pt0[kbyte]] ^ sbox[i ^ pt1[kbyte]]) == (glitched_ct0[kbyte] ^ glitched_ct1[kbyte]):\n",
    "            if (sbox[i ^ pt0[kbyte]] ^ sbox[i ^ pt2[kbyte]]) == (glitched_ct0[kbyte] ^ glitched_ct2[kbyte]):\n",
    "                print(i)\n",
    "                key_guess.append(i)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bytearray(key_guess)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You'll see that you've got the correct key bytes, but they're out of order! This is because we didn't account for the ShiftRows operation (this doesn't matter for the analysis since we were using repeated bytes for the plaintext. If this wasn't the case, we'd have to match the plaintext up with the ciphertext instead of just leaving it to the end). If we swap the bytes of the key around based on ShiftRows:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "SR = [0, 13, 10, 7, 4, 1, 14, 11, 8, 5, 2, 15, 12, 9, 6, 3]\n",
    "for i in range(16):\n",
    "    print(hex(key_guess[SR[i]]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can print the key in the correct order.\n",
    "\n",
    "## Conclusions & Next Steps\n",
    "\n",
    "Though we didn't glitch an unmodified version of TINYAES128C, we did see that, given the right implementation of AES, we can use glitching to skip the repeated rounds of AES, thus bypassing its security and allowing us to recover the key. In the next lab, we'll look at mounting an attack against an unmodified TINYAES."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "<small>NO-FUN DISCLAIMER: This material is Copyright (C) NewAE Technology Inc., 2015-2020. ChipWhisperer is a trademark of NewAE Technology Inc., claimed in all jurisdictions, and registered in at least the United States of America, European Union, and Peoples Republic of China.\n",
    "\n",
    "Tutorials derived from our open-source work must be released under the associated open-source license, and notice of the source must be *clearly displayed*. Only original copyright holders may license or authorize other distribution - while NewAE Technology Inc. holds the copyright for many tutorials, the github repository includes community contributions which we cannot license under special terms and **must** be maintained as an open-source release. Please contact us for special permissions (where possible).\n",
    "\n",
    "THE SOFTWARE IS PROVIDED \"AS IS\", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.</small>"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
